import gdal
import osr
import gdalnumeric
import calendar
import numpy as np
import netCDF4 as nc
import xarray as xr
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import RasterHelper
import datetime
import os
import nc2tiff
import MPSD
import string
from multiprocessing import Process

def extracteIndex(in_root,out_root,varibles=['ssrd','sund','tcc'],years=['2002','2015']):
	gdal.AllRegister()
	for year in range(int(years[0]),int(years[1])+1):
		print('**************processing %d******************'%year)
		for varible in varibles:
			print('**************processing %s******************'%varible)
			file_root=os.path.join(in_root,'%d'%year,varible)
			# if varible=='tcc':
			# 	info_ds=gdal.Open(os.path.join(file_root,'%s_%d01010010000.tif'%(varible,year)))
			# else:
			info_ds=gdal.Open(os.path.join(file_root,'%s_%d0101001.tif'%(varible,year)))
			geotiff_info=nc2tiff.getSubDatasetInfo(info_ds)
			geotiff_info.append([info_ds.RasterXSize,info_ds.RasterYSize])
			info_ds=None
			year_data=np.empty((12,geotiff_info[-1][1],geotiff_info[-1][0]))
			if varible == 'sund':
				year_valid_days=np.empty((12,geotiff_info[-1][1],geotiff_info[-1][0]))
			# month iterater
			for month in range(1,13):
				days=calendar.monthrange(year, month)[1]
				month_data=np.empty((days,geotiff_info[-1][1],geotiff_info[-1][0]))

				for day in range(1,days+1):

					cur_date=datetime.datetime.strptime('%d%02d%d'%(year,month,day),'%Y%m%d')
					# if varible=='tcc':
					# 	file_name='%s_%s.tif'%(varible,datetime.datetime.strftime(cur_date,'%Y%m%d%j%H%M'))
					# else:
					file_name='%s_%s.tif'%(varible,datetime.datetime.strftime(cur_date,'%Y%m%d%j'))
					print('==========read %s==========='%file_name)
					cur_ds=gdal.Open(os.path.join(file_root,file_name))

					month_data[day-1]=cur_ds.GetRasterBand(1).ReadAsArray()

					cur_ds =None
				nc2tiff.isExists(os.path.join(out_root,varible))

				if varible =='tcc':
					year_data[month-1]=np.mean(month_data,axis=0)
					# monthly average total cloud cover
					RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%d%02d_sum.tif'%(varible,year,month)),year_data[month-1],geotiff_info[0],geotiff_info[1],geotiff_info[2])
				else:
					# monthly total sunshine duration and total surface shortwave radiation downward
					year_data[month-1]=np.sum(month_data,axis=0)
					RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%02d%02d_sum.tif'%(varible,year,month)),year_data[month-1],geotiff_info[0],geotiff_info[1],geotiff_info[2])
					if varible=='sund':
						# monthly vaild sunshine days
						year_valid_days[month-1]=np.sum(month_data>=6*3600,axis=0)
						RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%d%02d_valid_sd_days.tif'%(varible,year,month)),year_valid_days[month-1],geotiff_info[0],geotiff_info[1],geotiff_info[2])

					#if
				#if-else
				month_data=None
				percentage =(((2006-float(years[0]))*12+6)/((float(years[1])-float(years[0]))*12.0))*100
				print('------------processing %.2f ---------'%percentage)
			# for month

			# yearly total sunshine duration and total surface shortwave radiation downward and average total cloud cover
			if varible == 'tcc':
				RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%d_yearly.tif'%(varible,year)),np.mean(year_data,axis=0),geotiff_info[0],geotiff_info[1],geotiff_info[2])
			else:
				RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%d_yearly.tif'%(varible,year)),np.sum(year_data,axis=0),geotiff_info[0],geotiff_info[1],geotiff_info[2])
				if varible=='sund':
					RasterHelper.writeGeoTiff(os.path.join(out_root,varible,'%s_%d_yearly_valid_sd_days.tif'%(varible,year)),np.sum(year_valid_days,axis=0),geotiff_info[0],geotiff_info[1],geotiff_info[2])
					year_valid_days=None
			year_data=None

if __name__ == '__main__':
	p1=Process(target=extracteIndex,args=('/run/media/hugo/Hugo Zhang/ERA','/run/media/hugo/Hugo Data/Result',),kwargs={'years':['2000','2015'],'varibles':['tcc']})
	# p2=Process(target=extracteIndex,args=('/run/media/hugo/Hugo Zhang/ERA','/run/media/hugo/Hugo Data/Result',),kwargs={'years':['2002','2015'],'varibles':['tcc']})
	# p2=Process(target=extracteIndex,args=('/run/media/hugo/Hugo Data/Result/','/run/media/hugo/Hugo Zhang/ERA_P/Index'),kwargs={'years':['2013','2015']})
	p1.start()
	# p2.start()


